#include "compiler/lexer.h"
#include <cassert>
#include <iostream>
#include <string>
#include <vector>

// Helper function to print token type
std::string tokenTypeToString(TokenType type) {
    switch (type) {
        case TokenType::KEYWORD: return "KEYWORD";
        case TokenType::IDENTIFIER: return "IDENTIFIER";
        case TokenType::CONST_INT: return "CONST_INT";
        case TokenType::CONST_STRING: return "CONST_STRING";
        case TokenType::OPERATOR: return "OPERATOR";
        case TokenType::DELIMITER: return "DELIMITER";
        case TokenType::END: return "END";
        case TokenType::ERROR: return "ERROR";
        default: return "UNKNOWN";
    }
}

// Helper function to compare tokens
bool tokensEqual(const Token& token, TokenType expectedType, const std::string& expectedValue) {
    return token.type == expectedType && token.value == expectedValue;
}

// Test error handling functionality
void testErrorHandling() {
    std::cout << "Running error handling test..." << std::endl;
    
    std::string sql = "'unclosed string";
    Lexer lexer(sql);
    bool exceptionCaught = false;
    
    try {
        lexer.tokenize();
    } catch (const LexerError& e) {
        exceptionCaught = true;
        std::cout << "Expected error caught: " << e.what() << std::endl;
    }
    
    assert(exceptionCaught);
    std::cout << "Error handling test passed!" << std::endl;
}

// Test keyword recognition
void testKeywordRecognition() {
    std::cout << "Running keyword recognition test..." << std::endl;
    
    std::string sql = "SELECT FROM WHERE CREATE TABLE INSERT INTO VALUES DELETE INT VARCHAR AND OR BETWEEN";
    Lexer lexer(sql);
    std::vector<Token> tokens = lexer.tokenize();
    
    // Expect 14 keywords plus END token
    assert(tokens.size() == 15);
    
    // Check that all are recognized as keywords
    for (size_t i = 0; i < tokens.size() - 1; ++i) {
        assert(tokens[i].type == TokenType::KEYWORD);
    }
    
    // Check END token
    assert(tokens.back().type == TokenType::END);
    
    std::cout << "Keyword recognition test passed!" << std::endl;
}

// Test identifier recognition
void testIdentifierRecognition() {
    std::cout << "Running identifier recognition test..." << std::endl;
    
    std::string sql = "table1 column_name _user _123 test_123ABC";
    Lexer lexer(sql);
    std::vector<Token> tokens = lexer.tokenize();
    
    // Expect 5 identifiers plus END token
    assert(tokens.size() == 6);
    
    // Check that all are recognized as identifiers
    for (size_t i = 0; i < tokens.size() - 1; ++i) {
        assert(tokens[i].type == TokenType::IDENTIFIER);
    }
    
    // Check END token
    assert(tokens.back().type == TokenType::END);
    
    std::cout << "Identifier recognition test passed!" << std::endl;
}

// Test constant integer recognition
void testConstantIntegerRecognition() {
    std::cout << "Running constant integer recognition test..." << std::endl;
    
    std::string sql = "123 4567 0 987654321";
    Lexer lexer(sql);
    std::vector<Token> tokens = lexer.tokenize();
    
    // Expect 4 integer constants plus END token
    assert(tokens.size() == 5);
    
    // Check that all are recognized as integer constants
    for (size_t i = 0; i < tokens.size() - 1; ++i) {
        assert(tokens[i].type == TokenType::CONST_INT);
    }
    
    // Check specific values
    assert(tokens[0].value == "123");
    assert(tokens[1].value == "4567");
    assert(tokens[2].value == "0");
    assert(tokens[3].value == "987654321");
    
    // Check END token
    assert(tokens.back().type == TokenType::END);
    
    std::cout << "Constant integer recognition test passed!" << std::endl;
}

// Test string constant recognition
void testStringConstantRecognition() {
    std::cout << "Running string constant recognition test..." << std::endl;
    
    std::string sql = "'hello' 'world' 'test string' ''";
    Lexer lexer(sql);
    std::vector<Token> tokens = lexer.tokenize();
    
    // Expect 4 string constants plus END token
    assert(tokens.size() == 5);
    
    // Check that all are recognized as string constants
    for (size_t i = 0; i < tokens.size() - 1; ++i) {
        assert(tokens[i].type == TokenType::CONST_STRING);
    }
    
    // Check specific values
    assert(tokens[0].value == "hello");
    assert(tokens[1].value == "world");
    assert(tokens[2].value == "test string");
    assert(tokens[3].value == "");
    
    // Check END token
    assert(tokens.back().type == TokenType::END);
    
    std::cout << "String constant recognition test passed!" << std::endl;
}

// Test operator and delimiter recognition
void testOperatorAndDelimiterRecognition() {
    std::cout << "Running operator and delimiter recognition test..." << std::endl;
    
    std::string sql = "( ) , ; = == > >= < <= !=";
    Lexer lexer(sql);
    std::vector<Token> tokens = lexer.tokenize();
    
    // Expect 11 operators/delimiters plus END token
    assert(tokens.size() == 12);
    
    // Check specific types and values
    assert(tokensEqual(tokens[0], TokenType::DELIMITER, "("));
    assert(tokensEqual(tokens[1], TokenType::DELIMITER, ")"));
    assert(tokensEqual(tokens[2], TokenType::DELIMITER, ","));
    assert(tokensEqual(tokens[3], TokenType::DELIMITER, ";"));
    assert(tokensEqual(tokens[4], TokenType::OPERATOR, "="));
    assert(tokensEqual(tokens[5], TokenType::OPERATOR, "=="));
    assert(tokensEqual(tokens[6], TokenType::OPERATOR, ">"));
    assert(tokensEqual(tokens[7], TokenType::OPERATOR, ">="));
    assert(tokensEqual(tokens[8], TokenType::OPERATOR, "<"));
    assert(tokensEqual(tokens[9], TokenType::OPERATOR, "<="));
    assert(tokensEqual(tokens[10], TokenType::OPERATOR, "!="));
    
    // Check END token
    assert(tokens.back().type == TokenType::END);
    
    std::cout << "Operator and delimiter recognition test passed!" << std::endl;
}

// Test whitespace and comment skipping
void testWhitespaceAndCommentSkipping() {
    std::cout << "Running whitespace and comment skipping test..." << std::endl;
    
    std::string sql = "SELECT     -- This is a single line comment\nFROM table1";
    Lexer lexer(sql);
    std::vector<Token> tokens = lexer.tokenize();
    
    // Expect 3 tokens (SELECT, FROM, table1) plus END token
    assert(tokens.size() == 4);
    
    // Check specific types and values
    assert(tokensEqual(tokens[0], TokenType::KEYWORD, "SELECT"));
    assert(tokensEqual(tokens[1], TokenType::KEYWORD, "FROM"));
    assert(tokensEqual(tokens[2], TokenType::IDENTIFIER, "table1"));
    
    // Check END token
    assert(tokens.back().type == TokenType::END);
    
    std::cout << "Whitespace and comment skipping test passed!" << std::endl;
}

// Test complete SQL statement parsing
void testCompleteSQLStatement() {
    std::cout << "Running complete SQL statement test..." << std::endl;
    
    std::string sql = "SELECT id, name FROM users WHERE age > 18;";
    Lexer lexer(sql);
    std::vector<Token> tokens = lexer.tokenize();
    
    // Expect 11 tokens plus END token
    assert(tokens.size() == 12);
    
    // Check specific types and values
    assert(tokensEqual(tokens[0], TokenType::KEYWORD, "SELECT"));
    assert(tokensEqual(tokens[1], TokenType::IDENTIFIER, "id"));
    assert(tokensEqual(tokens[2], TokenType::DELIMITER, ","));
    assert(tokensEqual(tokens[3], TokenType::IDENTIFIER, "name"));
    assert(tokensEqual(tokens[4], TokenType::KEYWORD, "FROM"));
    assert(tokensEqual(tokens[5], TokenType::IDENTIFIER, "users"));
    assert(tokensEqual(tokens[6], TokenType::KEYWORD, "WHERE"));
    assert(tokensEqual(tokens[7], TokenType::IDENTIFIER, "age"));
    assert(tokensEqual(tokens[8], TokenType::OPERATOR, ">"));
    assert(tokensEqual(tokens[9], TokenType::CONST_INT, "18"));
    assert(tokensEqual(tokens[10], TokenType::DELIMITER, ";"));
    
    // Check END token
    assert(tokens.back().type == TokenType::END);
    
    std::cout << "Complete SQL statement test passed!" << std::endl;
}

int main() {
    std::cout << "Running Lexer unit tests..." << std::endl;
    std::cout << "==================================" << std::endl;
    
    // Run individual tests
    testErrorHandling();
    testKeywordRecognition();
    testIdentifierRecognition();
    testConstantIntegerRecognition();
    testStringConstantRecognition();
    testOperatorAndDelimiterRecognition();
    testWhitespaceAndCommentSkipping();
    testCompleteSQLStatement();
    
    std::cout << "==================================" << std::endl;
    std::cout << "All Lexer tests passed successfully!" << std::endl;
    return 0;
}